import os
from glob import glob
import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from src.utils import preprocess_data

diagnoses = ['abnormal', 'acl', 'meniscus']
mrType = ['axial', 'coronal', 'sagittal']


class MRNetDataset(Dataset):
    def __init__(self, dataset_dir, labels_path, choose_num=16, transform=None, device=None):
        self.case_paths = sorted(glob(f'{dataset_dir}/**.npy'))
        self.abnormal_df = pd.read_csv(os.path.join(labels_path, 'abnormal.csv'), header=None)
        self.acl_df = pd.read_csv(os.path.join(labels_path, 'acl.csv'), header=None)
        self.meniscus_df = pd.read_csv(os.path.join(labels_path, 'meniscus.csv'), header=None)
        self.transform = transform
        self.device = device
        if self.device is None:
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.choose_num = choose_num


    def __len__(self):
        return len(self.case_paths)

    def __getitem__(self, idx):
        case_path = self.case_paths[idx]
        series = preprocess_data(case_path, self.transform, choose_num=self.choose_num)
        # 对应的序号
        case_id = int(os.path.splitext(os.path.basename(case_path))[0])
        # 三个对应的诊断
        abnormal_diagnose_df = self.abnormal_df[self.abnormal_df.iloc[:, 0] == case_id]
        acl_diagnose_df = self.acl_df[self.acl_df.iloc[:, 0] == case_id]
        meniscus_diagnose_df = self.meniscus_df[self.meniscus_df.iloc[:, 0] == case_id]
        # 整合三个诊断
        diagnoses = [abnormal_diagnose_df.iloc[0, 1].astype(np.float32), acl_diagnose_df.iloc[0, 1].astype(np.float32),
                     meniscus_diagnose_df.iloc[0, 1].astype(np.float32)]
        labels = torch.tensor(diagnoses)
        labels = labels.repeat(16, 1)

        return series, labels


def make_dataset(data_dir, dataset_type, mr_type, device=None):
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

    dataset_dir = f'{data_dir}/{dataset_type}/{mr_type}'
    labels_path = f'{data_dir}/{dataset_type}'

    if dataset_type == 'train':
        transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.RandomHorizontalFlip(),
            transforms.RandomAffine(25, translate=(0.1, 0.1)),
            transforms.ToTensor(),
        ])

    elif dataset_type == 'valid':
        transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.ToTensor()
        ])
    else:
        raise ValueError('Dataset needs to be train or valid.')

    if mr_type not in mrType:
        raise ValueError('mr_type needs to be {} or {} or {}'.format(mrType[0], mrType[1], mrType[2]))

    dataset = MRNetDataset(dataset_dir, labels_path, choose_num=16, transform=transform, device=device)

    return dataset


def make_data_loader(data_dir, dataset_type, plane, device=None, shuffle=False):
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'

    dataset = make_dataset(data_dir, dataset_type, plane, device=device)
    data_loader = DataLoader(dataset, batch_size=None, shuffle=shuffle)

    return data_loader


if __name__ == '__main__':
    a = make_data_loader('../data/MRNet-v1.0', 'train', 'axial')
    for iteration, batch in enumerate(a):
        images, labels = batch
        print(images.size())
        print(labels)
